﻿using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.IO;
using System.Linq;
using System.Text;

namespace GPUTracer
{
    public class KDTreeConverter
    {
        private List<Object> _objects = new List<Object>();

        public byte[] ConvertKDTree(KDTreeElement kdtree, Object[] objects)
        {
            //First we have to add all the planes:
            foreach (var o in objects)
            {
                if (o is Plane)
                    _objects.Add(o);
            }
            //Then all the other objects:
            foreach (var o in objects)
            {
                if (!(o is Plane))
                    _objects.Add(o);
            }


            var memstream = new MemoryStream();
            var sceneStream = new BinaryWriter(memstream);

            var treeElementBytes = ConvertElement(kdtree);

            var header = new byte[4 * 4];
            var bytesSkipped = treeElementBytes.Count() + 4*4*3;
            var dummySpace = 16 - (bytesSkipped % 16) % 16;
            Buffer.BlockCopy(BitConverter.GetBytes(bytesSkipped + dummySpace), 0, header, 0, 4);

            var numPlanes = objects.Count(x => x is Plane);
            Buffer.BlockCopy(BitConverter.GetBytes(numPlanes), 0, header, 4, 4);

            var numObjects = objects.Count();
            Buffer.BlockCopy(BitConverter.GetBytes(numObjects), 0, header, 8, 4);

            var bboxHeader = new byte[4 * 8];
            int c = 0;
            Helper.Float4ToByteArray(bboxHeader, kdtree.BBox.Min, ref c);
            Helper.Float4ToByteArray(bboxHeader, kdtree.BBox.Max, ref c);

            sceneStream.Write(header);
            sceneStream.Write(bboxHeader);
            sceneStream.Write(treeElementBytes);

            var dummy = new byte[dummySpace];
            sceneStream.Write(dummy);

            foreach (var o in _objects)
            {
                sceneStream.Write(o.ToBinary());
            }

            return memstream.ToArray();
        }

        private byte[] ConvertElement(KDTreeElement element)
        {
            var memstream = new MemoryStream();
            var writer = new BinaryWriter(memstream);
            
            if (element is KDTreeLeaf)
            {
                int first;
                int second;
                first = element.NumObjects << 2;
                second = AddObjectsToList(((KDTreeLeaf)element).ContainedObjects);
                writer.Write(first);
                writer.Write(second);
            }
            else if (element is KDTreeNode)
            {
                float first = ((KDTreeNode) element).SplitPosition;
                var firstBytes = BitConverter.GetBytes(first);
                firstBytes[0] = (byte)((firstBytes[0] >> 2) << 2);
                switch (((KDTreeNode)element).SplitAxis)
                {
                    case Axis.XAxis:
                        firstBytes[0] = (byte) (firstBytes[0] | 1);
                        break;
                    case Axis.YAxis:
                        firstBytes[0] = (byte)(firstBytes[0] | 2);
                        break;
                    case Axis.ZAxis:
                        firstBytes[0] = (byte)(firstBytes[0] | 3);
                        break;
                    default:
                        throw new ArgumentOutOfRangeException();
                }

                BitConverter.ToSingle(firstBytes, 0);

                var firstChild = ConvertElement(((KDTreeNode) element).LeftChild);
                var secondChild = ConvertElement(((KDTreeNode)element).RightChild);

                int aboveChild = (firstChild.Count() / 8)+1;

                writer.Write(firstBytes);
                writer.Write(aboveChild);

                writer.Write(firstChild);
                writer.Write(secondChild);
            }

            Debug.Assert(memstream.Length % 8 == 0);
            return memstream.ToArray();
        }

        private int AddObjectsToList(List<Object> objects)
        {
            if (objects.Count == 0)
                return 0;

            //search, if set is already available:
            var searchObjSet = new HashSet<Object>(objects);
            int i = 0;
            bool c = false;
            while (i < _objects.Count)
            {
                if (searchObjSet.Count == 0)
                {
                    break;
                }

                if (searchObjSet.Contains(_objects[i]))
                {
                    searchObjSet.Remove(_objects[i]);
                    c = true;
                }
                else if (c)
                {
                    searchObjSet = new HashSet<Object>(objects);
                    c = false;
                    continue;
                }
                i++;
            }
            if (searchObjSet.Count == 0)
            {
                return i - objects.Count;
            }

            //If not, add it:
            int index = _objects.Count();
            foreach (var o in objects)
            {
                _objects.Add(o);
            }
            return index;
        }
    }
}
